from mootdx.reader import Reader
import pandas as pd
import os
import sys
import json
import time
import subprocess
import akshare as ak

# market 参数 std 为标准市场(就是股票), ext 为扩展市场(期货，黄金等)
# tdxdir 是通达信的数据目录, 根据自己的情况修改

# 请根据你的实际通达信安装路径修改下面这行
tdx_dir = 'C:/new_tdx'  # 修改为你的通达信安装目录

print(f"正在使用通达信数据目录: {tdx_dir}")
if not os.path.exists(tdx_dir):
    print(f"错误: 通达信目录 {tdx_dir} 不存在，请检查路径是否正确")
    sys.exit(1)

reader = Reader.factory(market='std', tdxdir=tdx_dir)

# 检查stock_list.json是否存在，如果存在则直接读取，否则重新获取
if os.path.exists('stock_list.json'):
    print("检测到stock_list.json文件，直接读取...")
    with open('stock_list.json', 'r', encoding='utf-8') as f:
        stock_list_data = json.load(f)
    # 从JSON数据创建DataFrame
    stock_list = stock_list_data["stocks"]
    filtered_stocks = pd.DataFrame(stock_list)
    print(f"从文件中读取到 {len(filtered_stocks)} 只股票")
else:
    # 使用akshare获取A股实时数据
    print("正在获取A股实时数据...")
    df = ak.stock_zh_a_spot_em()
    
    # 选择需要的列并重命名
    all_stocks = df[["代码", "名称", "总市值"]].rename(columns={"代码": "code", "名称": "name", "总市值": "market_cap"})
    
    # 补全股票代码为6位
    all_stocks["code"] = all_stocks["code"].str.zfill(6)
    
    # 按代码排序
    all_stocks = all_stocks.sort_values("code").reset_index(drop=True)
    
    print(f"总共找到 {len(all_stocks)} 只证券产品")

    # 筛选个股（排除创业板、科创板、北交所和指数）
    # 排除规则：
    # 1. 创业板：300、301开头
    # 2. 科创板：688开头
    # 3. 北交所：8、4开头
    # 4. 指数：名称中包含"指数"字样的证券
    # 5. 总市值大于0
    # 6. 退市股：名称中包含"退市"字样的证券
    # 7. ST股：名称中包含"ST"字样的证券
    filtered_stocks = all_stocks[
        ~all_stocks["code"].str.startswith(("300", "301", "688", "8", "4")) & 
        ~all_stocks["name"].str.contains("指数") &
        ~all_stocks["name"].str.contains("退市") &
        ~all_stocks["name"].str.contains("ST") &
        (all_stocks["market_cap"] > 0)
    ].copy()
    
    # 进一步筛选，只保留代码大于等于000300的证券，排除大部分指数
    # 但也保留以6开头的沪市股票
    filtered_stocks = filtered_stocks[
        (filtered_stocks["code"] >= "000300") | (filtered_stocks["code"].str.startswith("6"))
    ]
    
    print(f"符合条件的个股数量: {len(filtered_stocks)}")

    # 按市值从大到小排序
    filtered_stocks = filtered_stocks.sort_values("market_cap", ascending=False)
    
    # 去除重复代码（如果有的话）
    filtered_stocks = filtered_stocks.drop_duplicates(subset=['code'])
    
    print(f"去重后个股数量: {len(filtered_stocks)}")

    # 将股票列表保存到JSON文件
    stock_list_data = {
        "stocks": filtered_stocks[['code', 'name']].to_dict('records')
    }

    # 保存到JSON文件
    with open('stock_list.json', 'w', encoding='utf-8') as f:
        json.dump(stock_list_data, f, ensure_ascii=False, indent=2)

    print(f"股票列表已保存到 stock_list.json 文件中")

# 创建输出目录
output_dir = 'tdx_data'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 读取日线数据
count = 0
total = len(filtered_stocks)
start_time = time.time()
for index, stock in filtered_stocks.iterrows():
    stock_code = stock['code']
    stock_name = stock['name'].replace(' ', '')  # 清理股票名称中的空格
    count += 1
    print(f"[{count}/{total}] 读取{stock_code}日线数据...")
    
    try:
        daily_data = reader.daily(symbol=stock_code)
        if daily_data is not None and not daily_data.empty:
            # 重置索引，使日期成为普通列，然后保存到CSV
            # 保留两位小数
            daily_data_rounded = daily_data.reset_index()
            numeric_columns = ['open', 'high', 'low', 'close', 'amount', 'volume']
            for col in numeric_columns:
                if col in daily_data_rounded.columns:
                    daily_data_rounded[col] = daily_data_rounded[col].round(2)
            
            # 调整列的顺序为: date,open,close,high,low,volume,amount
            column_order = ['date', 'open', 'close', 'high', 'low', 'volume', 'amount']
            # 检查所有需要的列是否都存在
            available_columns = [col for col in column_order if col in daily_data_rounded.columns]
            daily_data_rounded = daily_data_rounded[available_columns]
            # 在文件名中包含股票名称
            daily_data_rounded.to_csv(f'{output_dir}/{stock_code}_{stock_name}_daily.csv', index=False)
            print(f"{stock_code}日线数据已保存到 {output_dir}/{stock_code}_{stock_name}_daily.csv")
        else:
            print(f"未能读取到{stock_code}日线数据，请确认数据是否存在")
    except Exception as e:
        print(f"读取{stock_code}日线数据时发生错误: {e}")

end_time = time.time()
print(f"\n数据读取和保存操作完成! 总耗时: {end_time - start_time:.2f}秒")

# 数据获取完成后，自动执行数据处理功能，如果处理出错可以注释掉下面代码
print("\n开始执行数据处理功能...")
try:
    result = subprocess.run(['python', 'TDX_data_process.py'], 
                          capture_output=True, text=True, 
                          cwd=os.getcwd())
    
    if result.returncode == 0:
        print("数据处理完成!")
        print(result.stdout)
    else:
        print("数据处理过程中出现错误:")
        print(result.stderr)
        
except Exception as e:
    print(f"执行数据处理时发生异常: {e}")